# %%
import torchvision
import torch

from cnn import *
# %%
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
train_data = torchvision.datasets.CIFAR10(
    root="data/cifar10", train=True, transform=torchvision.transforms.ToTensor(), download=True)
test_data = torchvision.datasets.CIFAR10(
    root="data/cifar10", train=False, transform=torchvision.transforms.ToTensor(), download=True)
train_loader = torch.utils.data.DataLoader(train_data, batch_size=64)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=64)
# %%
cnn = CNN().to(device)
# cnn=cnn.cuda()
loss_cross = nn.CrossEntropyLoss().to(device)

learning_rate = 1e-2
optimizer = torch.optim.SGD(cnn.parameters(), lr=learning_rate)

# %%
total_step = 0
total_step = 0
epoch = 10
# %%
for i in range(epoch):
    print("___________________{}________________".format(i+1))

    for data in train_loader:
        imgs, targets = data
        imgs = imgs.to(device)
        targets = targets.to(device)
        outputs = cnn(imgs)
        loss = loss_cross(outputs, targets)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_step += 1

        if total_step % 100 == 0:
            print("-------------------*******-----------------------------")
            print("loss:{}".format(loss.item()))
            print("total_step:{}".format(total_step))
            print("----------------------------------------------------")
            correct = 0
            total = 0
            with torch.no_grad():
                for data in test_loader:
                    imgs, targets = data
                    imgs = imgs.to(device)
                    targets = targets.to(device)
                    outputs = cnn(imgs)
                    _, predicted = torch.max(outputs.data, 1)
                    total += targets.size(0)
                    correct += (predicted == targets).sum().item()
            print("correct:{}".format(correct))
            print("total:{}".format(total))
            print("accuracy:{}".format(correct/total))
